library(dplyr)
library(tidyr)
library(ggplot2)
library(viridis)
library(ggpubr)
library(patchwork)
library(circlize)
library(scales)

if (!file.exists("output")) {
  dir.create("output")
}

# load data
data <- read.csv("Figure_2a_data.csv")

# Extract signatures with significant differences by subsite
# Significance is assessed using a two-sided Kruskal-Wallis test and Bonferroni correction.
signatures <- names(data[, -c(1:2)])
stat_all <- NULL
for (s in signatures) {
  kruskal_p <- kruskal.test(data[, s], data[, "subsite"])$p.value
  stat <- data.frame(
    signature = paste(s),
    pval = kruskal_p
  ) %>%
    mutate(p_adj = ifelse(pval * length(signatures) < 1, pval * length(signatures), 1))
  stat_all <- rbind(stat_all, stat)
}
signif <- stat_all %>%
  filter(p_adj < 0.05) %>%
  pull(signature)

data <- data %>% mutate(subsite = factor(subsite, levels = c("OC", "OPC", "Hypopharynx", "Larynx")))

# Define function to generate bubble plot
# The size of each dot represents the proportion of samples presenting each mutational signature in the whole HNC dataset and across anatomical subsites.
# The color represents the mean relative attribution of each signature.
# Gray dots indicate signatures without significantly different relative burdens by subsite.

# sigs: Data frame with signature attributions
# significant_sigs: Character vector containing signature names with significantly different relative activities across aatomical subsites

sig_landscape <- function(sigs, significant_sigs) {
  signatures <- names(sigs)[-c(1:2)]
  # calculate frequency and mean relative attribution in tumors from each signature by subsite
  sigs_rel <- sigs[-2]
  sigs_cat <- sigs_rel %>%
    tibble::column_to_rownames("donor_id") %>%
    mutate_if(is.numeric, ~ ifelse(. > 0, 1, 0))
  perc_sig <- sigs_rel %>%
    summarise_if(is.numeric, mean) %>%
    t()

  summary <- sigs_cat %>%
    summarise_if(is.numeric, sum) %>%
    t() %>%
    as.data.frame() %>%
    rename(n_pos = "V1") %>%
    mutate(freq = n_pos / nrow(sigs)) %>%
    cbind(perc_sig) %>%
    tibble::rownames_to_column("sig")

  order <- summary[order(1:nrow(summary)), ]

  # calculate frequency and mean relative attribution in tumors from each signature by subsite
  summary_bysubsite <- NULL
  for (sub in levels(sigs$subsite)) {
    sigs_cat.sub <- sigs_cat[sigs[which(sigs$subsite == sub), "donor_id"], ]
    perc_sig.sub <- sigs_rel %>%
      filter(donor_id %in% rownames(sigs_cat.sub)) %>%
      summarise_if(is.numeric, mean) %>%
      t()

    summary.sub <- sigs_cat.sub %>%
      summarise_if(is.numeric, sum) %>%
      t() %>%
      as.data.frame() %>%
      rename(n_pos = "V1") %>%
      mutate(freq = n_pos / nrow(sigs_cat.sub), subsite = sub) %>%
      cbind(perc_sig.sub) %>%
      tibble::rownames_to_column("sig")

    summary_bysubsite <- rbind(summary_bysubsite, summary.sub)
  }

  p_axis <- ggplot(summary, aes(x = sig)) +
    scale_x_discrete(name = "Signatures", limits = order$sig, position = "top") +
    theme_minimal() +
    theme(
      axis.title = element_blank(),
      panel.grid = element_blank(),
      axis.text.x = element_text(angle = 45, hjust = 0),
      axis.text.y = element_blank(),
      plot.margin = unit(c(0, 0.5, 0, 0.5), "lines")
    )
  # plot general
  p_all <- ggplot(summary, aes(y = "HNC", x = sig, color = perc_sig, size = ifelse(freq == 0, NA, freq))) +
    geom_point(shape = 19, stroke = 0.75) +
    scale_size(breaks = c(0.25, 0.50, 0.75, 1.00), limits = c(0, 1)) +
    scale_y_discrete(name = " ") +
    scale_x_discrete(name = "Signatures", limits = order$sig) +
    labs(
      color = "Mean relative \n attribution",
      size = "Proportion of tumors \n with the signature"
    ) +
    scale_color_viridis(limits = c(0, max(c(summary$perc_sig, summary_bysubsite$perc_sig.sub)))) +
    theme_minimal() +
    theme(
      legend.position = "none",
      panel.border = element_rect(colour = "black", fill = NA, size = .25),
      panel.grid.major = element_blank(),
      axis.title.x = element_blank(),
      axis.text.x = element_blank(),
      plot.margin = unit(c(0, 0.5, 0, 0.5), "lines")
    )

  # plot subsites

  summary_bysubsite <- summary_bysubsite %>%
    mutate(sig_cutoff = ifelse(sig %in% significant_sigs, "p adj < 0.05", "p adj > 0.05"))

  p_sub <- ggplot(summary_bysubsite, aes(
    x = sig, y = subsite,
    color = perc_sig.sub,
    size = ifelse(freq == 0, NA, freq),
    shape = sig_cutoff
  )) +
    geom_point(stroke = 0.75) +
    scale_x_discrete(name = "Signatures", limits = order$sig) +
    scale_y_discrete(name = " ", limits = rev(levels(data$subsite))) +
    scale_color_viridis(limits = c(0, max(c(summary$perc_sig, summary_bysubsite$perc_sig.sub)))) +
    scale_size(breaks = c(0.25, 0.50, 0.75, 1.00), limits = c(0, 1)) +
    scale_shape_manual(values = c(19, 21)) +
    labs(
      color = "Mean relative \n attribution",
      size = "Proportion of tumors \n with the signature",
      shape = ""
    ) +
    theme_minimal() +
    theme(
      legend.position = "bottom",
      panel.border = element_rect(colour = "black", fill = NA, size = .25),
      panel.grid.major = element_blank(),
      axis.title.x = element_blank(),
      axis.text.x = element_blank(),
      plot.margin = unit(c(0, 0.5, 0, 0.5), "lines")
    )

  wrap_plots(p_axis, p_all, p_sub, ncol = 1, heights = c(.1, 0.15, 0.5))
}

# generat plots for SBS, DBS, and ID signatures
for (s in list("SBS", "DBS", "ID")) {
  sigs <- data %>% select(c(1, 2, grep(s, names(data))))
  p <- sig_landscape(sigs, signif)
  print(p)
  ggsave(paste0("output/Figure_2a_", s, "_", Sys.Date(), ".pdf"),
    device = "pdf",
    width = 4.5, height = 5
  )
}

# TOP PANEL - TMB PLOT

source("HNC_metadata_tidy.R")

# Note: Mutational burdens and signature attributions can be found in the accompanying Supplementary Tables of the HNC manuscript
sbsCOSMIC <- read.csv("Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  mutate(`SBS7a-c` = SBS7a + SBS7b + SBS7c, .keep = "unused", .after = "SBS5") %>%
  mutate(`SBS17a-b` = SBS17a + SBS17b, .keep = "unused", .after = "SBS16") %>%
  tibble::rownames_to_column("donor_id")

dbsCOSMIC <- read.csv("Input_DBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  tibble::rownames_to_column("donor_id")

idCOSMIC <- read.csv("Input_ID_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  tibble::rownames_to_column("donor_id")

additional_df <- list(sbsCOSMIC, dbsCOSMIC, idCOSMIC)

cosmic <- additional_df[[1]]
for (n in 2:length(additional_df)) {
  cosmic <- cosmic %>% left_join(additional_df[[n]], by = "donor_id")
}

data <- data %>% left_join(cosmic, by = "donor_id")

for (sig_context in c("sbs", "dbs", "id")) {
  sig <- get(paste0(sig_context, "COSMIC"))
  ## order
  order <- sig[1]
  for (s in names(sig)[-1]) {
    df <- sig[, c("donor_id", s)] %>%
      filter(get(s) != 0) %>%
      arrange(get(s)) %>%
      tibble::rownames_to_column("order") %>%
      select("donor_id", "order", -s)
    names(df) <- c("donor_id", s)
    order <- left_join(order, df, by = "donor_id")
  }

  order <- order %>%
    gather(key = "signature", "order", -c("donor_id")) %>%
    mutate(code = paste(donor_id, signature, sep = "_"), .keep = "unused")

  sbs_perMb <- sig %>%
    mutate_if(is.numeric, \(x){
      x / 2800
    }) %>%
    gather(key = "signature", "mut_per_Mb", -c("donor_id")) %>%
    mutate(code = paste(donor_id, signature, sep = "_")) %>%
    left_join(order, by = "code") %>%
    filter(mut_per_Mb != 0) %>%
    mutate(order = as.numeric(order))

  median <- sbs_perMb %>%
    group_by(signature) %>%
    summarise(median_sig = median(mut_per_Mb)) %>%
    arrange(desc(median_sig))

  sbs_perMb <- sbs_perMb %>% mutate(signature = factor(signature, levels = median$signature))

  TMBplot <- ggplot(sbs_perMb, aes(x = order, y = mut_per_Mb)) +
    facet_grid(. ~ factor(signature, levels = median$signature),
      scales = "free_x"
    ) +
    geom_point(alpha = .3, size = .7, stroke = 0) +
    geom_hline(data = median, aes(yintercept = median_sig), size = .5, color = "red") +
    scale_y_continuous(trans = log10_trans()) +
    labs(y = "Mutations \n per megabase") +
    theme_minimal() +
    theme(
      panel.grid.major = element_blank(),
      panel.grid.minor.x = element_blank(),
      panel.spacing = unit(.1, "lines"),
      axis.text.y = element_text(size = 6),
      axis.text.x = element_blank(),
      axis.title = element_blank(),
      strip.text = element_text(angle = 45, size = 6),
      plot.margin = unit(c(0, 0.5, 0, 0.5), "lines")
    )

  ggsave(paste0("output/Figure_2a_TMBplot_", sig_context, "_", Sys.Date(), ".pdf"),
    device = "pdf",
    width = 4.5, height = 2
  )
}
